import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

from python_ai.common.xcommon import *
import tensorflow.compat.v1 as tf
import tensorflow as tsf
import numpy as np

np.random.seed(777)
tf.set_random_seed(777)

# Teach hello: hihell -> ihello

n_steps = 6
n_cls = 5
n_input = n_cls
n_hidden = 8
alpha = 0.1

idx2char = np.array(list('hielo'))  # (5, )
char2idx = dict()
for i, c in enumerate(idx2char):
    char2idx[c] = i

sep('Input x_data and y_data')
batch_size = 3
# x_data
x_data_sent = 'hihell' * batch_size
x_data_idx = [char2idx[c] for c in x_data_sent]
x_data_idx = np.array(x_data_idx).reshape(-1, n_steps)
check_shape(x_data_idx, 'x_data_idx')  # (m, 6)
x_data = np.eye(n_cls)[x_data_idx]
check_shape(x_data, 'x_data')  # (m, 6, 5)
# y_data
y_data_sent = 'ihello' * batch_size
y_data_idx = np.array([char2idx[c] for c in y_data_sent]).reshape(-1, n_steps)
check_shape(y_data_idx, 'y_data_idx')  # (m, 6)

sep('Graph')
ph_x = tf.placeholder(tf.float32, [None, n_steps, n_input], 'ph_x')
check_shape(ph_x, 'ph_x')  # (?, 6, 5)
ph_y = tf.placeholder(tf.int32, [None, n_steps], 'ph_y')
check_shape(ph_y, 'ph_y')  # (?, 6)
n_samples = tf.shape(ph_x)[0]

cell = tf.nn.rnn_cell.LSTMCell(n_hidden)
init_stat = cell.zero_state(n_samples, tf.float32)
check_shape(init_stat, 'init_stat')  # (2, ) # LSTMCell's state is a 2 elements tuple
outputs, states = tf.nn.dynamic_rnn(cell, ph_x, initial_state=init_stat)
check_shape(outputs, 'outputs')  # (m, 6, 8)

outputs_reshaped = tf.reshape(outputs, [-1, n_hidden])
check_shape(outputs_reshaped, 'outputs_reshaped')  # (6m, 8)
fc_out = tsf.contrib.layers.fully_connected(outputs_reshaped,
                                           n_cls,
                                           activation_fn=None)
check_shape(fc_out, 'fc_out')  # (6m, 5)
logits = tf.reshape(fc_out, [-1, n_steps, n_cls])
check_shape(logits, 'logits')  # (m, 6, 5)

weights = tf.ones([n_samples, n_steps], dtype=tf.float32)
check_shape(weights, 'weights')  # (m, 6)
cost = tsf.contrib.seq2seq.sequence_loss(logits=logits,
                                 targets=ph_y,
                                 weights=weights)
train = tf.train.AdamOptimizer(learning_rate=alpha).minimize(cost)
predict = tf.cast(tf.argmax(logits, axis=2), dtype=tf.int32)
check_shape(predict, 'predict')  # (m, 6)
acc = tf.reduce_mean(
    tf.cast(
        tf.equal(predict, ph_y),
        tf.float32
    )
)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    for i in range(20):
        feed_dict = {ph_x: x_data, ph_y: y_data_idx}
        _, costv, accv, predict_v = \
            sess.run([train, cost, acc, predict], feed_dict=feed_dict)

        print(f'#{i + 1}: cost = {costv}, acc = {accv}')
        print(x_data_sent)
        print(y_data_sent)
        print(''.join(idx2char[predict_v].ravel()))

        if np.isclose(1., accv):
            break
